import os
import argparse
import csv
from typing import List


def filter_subject_round(
    base_dir: str,
    src_round: int,
    dst_round: int,
    subjects: List[str],
):
    base = os.path.abspath(base_dir)
    src = os.path.join(base, f"round_{src_round}")
    dst = os.path.join(base, f"round_{dst_round}")
    os.makedirs(dst, exist_ok=True)
    names = ["train.csv", "val.csv", "test.csv"]
    subs_set = {s.strip().lower() for s in subjects}

    for name in names:
        src_csv = os.path.join(src, name)
        dst_csv = os.path.join(dst, name)
        if not os.path.isfile(src_csv):
            raise FileNotFoundError(src_csv)
        kept = 0
        with open(src_csv, "r", encoding="utf-8") as fsrc, open(
            dst_csv, "w", newline="", encoding="utf-8"
        ) as fdst:
            r = csv.reader(fsrc)
            w = csv.writer(fdst)
            header = next(r)
            w.writerow(header)
            try:
                subj_idx = header.index("subject")
            except ValueError:
                raise RuntimeError(f"'subject' 列不存在于 {src_csv}")
            for row in r:
                if len(row) <= subj_idx:
                    continue
                subj = row[subj_idx].strip().lower()
                if subj in subs_set:
                    w.writerow(row)
                    kept += 1
        print(f"{name}: kept={kept}, out={dst_csv}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="过滤指定 subject 的 train/val/test 到新的 round 目录"
    )
    parser.add_argument(
        "--split_base",
        type=str,
        required=False,
        default=r"E:\DEAP_pre\deap_face\subject-dependent\train-val-test",
        help="train/val/test 所在的 base 目录（包含 round_xx 子目录）",
    )
    parser.add_argument("--src_round", type=int, default=1, help="源 round 编号")
    parser.add_argument("--dst_round", type=int, default=101, help="目标 round 编号")
    parser.add_argument(
        "--subjects",
        type=str,
        nargs="+",
        default=["s01"],
        help="要保留的 subject 列表，如 s01 s02",
    )
    args = parser.parse_args()
    filter_subject_round(args.split_base, args.src_round, args.dst_round, args.subjects)
